import cv2
import numpy
import os
from umeyama import umeyama
import random

def get_image_paths(path):
    return [os.path.join(path, file) for file in os.listdir(path) if file.endswith(".jpg") or file.endswith(".png")]


def load_images( image_paths, convert=None ):
    iter_all_images = ( cv2.imread(fn) for fn in image_paths )
    if convert:
        iter_all_images = ( convert(img) for img in iter_all_images )
    for i,image in enumerate( iter_all_images ):
        if i == 0:
            all_images = numpy.empty( ( len(image_paths), ) + image.shape, dtype=image.dtype )
        all_images[i] = image
    return all_images


def random_transform( image, rotation_range, zoom_range, shift_range, random_flip ):
    h,w = image.shape[0:2]        #获取图片高宽
    rotation = numpy.random.uniform( -rotation_range, rotation_range )
    scale = numpy.random.uniform( 1 - zoom_range, 1 + zoom_range )
    tx = numpy.random.uniform( -shift_range, shift_range ) * w
    ty = numpy.random.uniform( -shift_range, shift_range ) * h
    mat = cv2.getRotationMatrix2D( (w//2,h//2), rotation, scale )
    mat[:,2] += (tx,ty)
    result = cv2.warpAffine( image, mat, (w,h), borderMode=cv2.BORDER_REPLICATE )
    if numpy.random.random() < random_flip:
        result = result[:,::-1]   #一定概率翻转
    return result


#get pair of random warped images from aligened face image
def random_warp( image, resolution ):
    assert image.shape == (resolution, resolution, 3)
    r = resolution//128
    range_ = numpy.linspace( 64*r-40*r, 64*r+40*r, 5 )
    mapx = numpy.broadcast_to( range_, (5,5) )
    mapy = mapx.T

    mapx = mapx + numpy.random.normal( size=(5,5), scale=5 )
    mapy = mapy + numpy.random.normal( size=(5,5), scale=5 )

    interp_mapx = cv2.resize( mapx, (40*r,40*r) )[4*r:36*r, 4*r:36*r].astype('float32')
    interp_mapy = cv2.resize( mapy, (40*r,40*r) )[4*r:36*r, 4*r:36*r].astype('float32')

    warped_image = cv2.remap( image, interp_mapx, interp_mapy, cv2.INTER_LINEAR )

    src_points = numpy.stack( [ mapx.ravel(), mapy.ravel() ], axis=-1 )
    dst_points = numpy.mgrid[0:32*r+1:8*r, 0:32*r+1:8*r].T.reshape(-1, 2)
    
    mat = umeyama( src_points, dst_points, True )[0:2]
    target_image = cv2.warpAffine( image, mat, (32*r, 32*r) )

    return warped_image, target_image


def get_training_data( images, batch_size, resolution ):
    random_transform_args = {
    'rotation_range': 10,
    'zoom_range': 0.05,
    'shift_range': 0.05,
    'random_flip': 0.4,
    }
    indices = numpy.random.randint( len(images), size=batch_size )
    for i,index in enumerate(indices):
        image = images[index]
        image = random_transform( image, **random_transform_args )
        warped_img, target_img = random_warp( image, resolution )

        if i == 0:
            warped_images = numpy.empty( (batch_size,) + warped_img.shape, warped_img.dtype )
            target_images = numpy.empty( (batch_size,) + target_img.shape, warped_img.dtype )

        warped_images[i] = warped_img
        target_images[i] = target_img
    return warped_images, target_images

def get_ratio_data( images, batch_size, resolution, idx_n ):
    random_transform_args = {
    'rotation_range': 10,
    'zoom_range': 0.05,
    'shift_range': 0.05,
    'random_flip': 0.4,
    }
    indices = numpy.random.randint( len(images), size=batch_size )
    for i,index in enumerate(indices):
        image = images[index]
        image = random_transform( image, **random_transform_args )
        warped_img, target_img = random_warp( image, resolution )

        if i == 0:
            warped_images = numpy.empty( (batch_size,) + warped_img.shape, warped_img.dtype )
            target_images = numpy.empty( (batch_size,) + target_img.shape, warped_img.dtype )
            noise_flag = []

        warped_images[i] = warped_img
        target_images[i] = target_img
        if index in idx_n:
            noise_flag.append(True)
        else:
            noise_flag.append(False)

    return warped_images, target_images, noise_flag